# LangChain的函数，工具和代理(六)：Conversational agent

import os
import time

import langchain_core.utils.function_calling as utils
import openai
from langchain.agents.format_scratchpad import format_to_openai_functions
from langchain.agents.output_parsers import OpenAIFunctionsAgentOutputParser
from langchain_core.prompts import MessagesPlaceholder
from pydantic import Field
from pydantic.v1 import BaseModel
# 导入langchain的tool
from langchain.tools import tool
import requests
import datetime
import wikipedia
from langchain.tools.render import format_tool_to_openai_function
from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.schema.agent import AgentFinish
from langchain.schema.runnable import RunnablePassthrough
from langchain.agents import AgentExecutor
from langchain.memory import ConversationBufferMemory

openai.api_key = "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"
os.environ['OPENAI_API_KEY'] = "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"


# 一、Tools
# 下面我们来定义一个函数用来查询天气的函数search
# 1、创建一个pydantic类
class OpenMeteoInput(BaseModel):
    latitude: float = Field(..., description="Latitude of the location to fetch weather data for")
    longitude: float = Field(..., description="Longitude of the location to fetch weather data for")


# 2、添加tool装饰器 args_schema是对输入参数做一个类型限制
# @tool(args_schema=OpenMeteoInput)
@tool
def get_current_temperature(latitude: float, longitude: float) -> str:
    """Fetch current temperature for given coordinates."""
    BASE_URL = "https://api.open-meteo.com/v1/forecast"

    # 请求的参数
    params = {
        'latitude': latitude,
        'longitude': longitude,
        'hourly': 'temperature_2m',
        'forecast_days': 1,
    }

    # 发起请求
    response = requests.get(BASE_URL, params=params)
    # 处理请求结果
    if response.status_code == 200:
        results = response.json()
    else:
        raise Exception(f"API Request failed with status code: {response.status_code}")

    current_utc_time = datetime.datetime.now()
    time_list = [datetime.datetime.fromisoformat(time_str.replace('Z', '+00:00')) for time_str in
                 results['hourly']['time']]
    temperature_list = results['hourly']['temperature_2m']

    closest_time_index = min(range(len(time_list)), key=lambda i: abs(time_list[i] - current_utc_time))
    current_temperature = temperature_list[closest_time_index]

    return f'The current temperature is {current_temperature}°C'


# 下面我们来定义一个函数用来查询维基百科
@tool
def search_wikipedia(query: str) -> str:
    """Run Wikipedia search and get page summaries."""
    page_titles = wikipedia.search(query)
    summaries = []
    for page_title in page_titles[: 3]:
        try:
            wiki_page = wikipedia.page(title=page_title, auto_suggest=False)
            summaries.append(f"Page: {page_title}\nSummary: {wiki_page.summary}")
        except (
                print("错误")
        ):
            pass
    if not summaries:
        return "No good Wikipedia Search Result was found"
    return "\n\n".join(summaries)


# 二、Routing
# 1、使用langchain的format_tool_to_openai_function方法将函数转换成描述变量比
wikipedia_function = utils.convert_to_openai_function(search_wikipedia)
temperature_function = utils.convert_to_openai_function(get_current_temperature)
functions = [temperature_function, wikipedia_function]
# 2、 接下来我们来定义一个可以查询天气温度和维基百科的chain:
# 2.1 创建函数描述变量
functions = [
    utils.convert_to_openai_function(f) for f in [
        search_wikipedia, get_current_temperature
    ]
]
# 2.2 创建prompt
prompt = ChatPromptTemplate.from_messages([
    ("system", "You are helpful but sassy assistant"),
    ("user", "{input}"),
])
# 2.3 定义llm
model = ChatOpenAI(temperature=0).bind(functions=functions)

# 4、定义chain
chain = prompt | model | OpenAIFunctionsAgentOutputParser()


# 5、创建一个agent的迭代流程
# 这里我们定义了一个run_agent的函数，该函数接受用户的输入信息，然后其内部会进行chain的迭代调用，直到chain返回的AgentFinish类型的消息则退出迭代。下面我们来测试一下该函数：
def run_agent(user_input):
    intermediate_steps = []
    while True:
        result = chain.invoke({
            "input": user_input,
            "agent_scratchpad": format_to_openai_functions(intermediate_steps)
        })
        if isinstance(result, AgentFinish):
            return result
        tool = {
            "search_wikipedia": search_wikipedia,
            "get_current_temperature": get_current_temperature,
        }
        observation = tool[result.tool].run(result.tool_input)
        intermediate_steps.append((result, observation))
        print("睡眠")
        time.sleep(30)
# 函数调用
# print(run_agent("上海现在的天气怎么样？"))



# 6、再次优化上面4和5的代码
# 这里我们在原来chain的基础上创建了一个agent_chain, 这个agent_chain会生成一个agent_scratchpad变量，它对应于prompt模板中的agent_scratchpad变量，并且在run_agent函数中我们不再调用chain变量而是调用agent_chain，同时在run_agent函数中我们将原先的agent_scratchpad变量换成了intermediate_steps，该变量用来存放中间结果如：intermediate_steps.append((result, observation))，经过这样的改造run_agent变得更加的简洁和优雅，可读性更强。
# rp = RunnablePassthrough.assign(
#     agent_scratchpad=lambda x: format_to_openai_functions(x["intermediate_steps"])
# )
# # 创建agent_chain
# agent_chain = rp | chain


# 创建一个agent的迭代流程
# def run_agent(user_input):
#     intermediate_steps = []
#     while True:
#         result = agent_chain.invoke({
#             "input": user_input,
#             "intermediate_steps": intermediate_steps
#         })
#         if isinstance(result, AgentFinish):
#             return result
#         tool = {
#             "search_wikipedia": search_wikipedia,
#             "get_current_temperature": get_current_temperature,
#         }
#         observation = tool[result.tool].run(result.tool_input)
#         intermediate_steps.append((result, observation))

# 到目前位置自动的优化的基础工作我们都已完成，接下来我们需要进一步完善agent的使用机制，不能每次都通过调用run_agent函数来实现和用户的互动，因此我们需要创建一个AgentExecutor，其内部已经封装了run_agent函数，使用起来更加方便：
# 创建AgentExecutor,打开调试器verbose=True
# agent_executor = AgentExecutor(agent=agent_chain, tools=[search_wikipedia,get_current_temperature], verbose=True)
# 执行agent_executor
# agent_executor.invoke({"input": "上海现在的天气怎么样?"})

# 7、再次优化上面6的代码 添加AgentExecutor+记忆组件
# 创建带有聊天历史记录变量的prompt模板
prompt = ChatPromptTemplate.from_messages([
    ("system", "You are helpful but sassy assistant"),
    MessagesPlaceholder(variable_name="chat_history"),
    ("user", "{input}"),
    MessagesPlaceholder(variable_name="agent_scratchpad")
])
# 创建agent_chain
agent_chain = RunnablePassthrough.assign(
    agent_scratchpad=lambda x: format_to_openai_functions(x["intermediate_steps"])
) | prompt | model | OpenAIFunctionsAgentOutputParser()

# 创建记忆力组件
memory = ConversationBufferMemory(return_messages=True,
                                  memory_key="chat_history")
# 创建AgentExecutor,打开调试器verbose=True，添加memory组件，添加tools工具组
agent_executor = AgentExecutor(agent=agent_chain,
                               tools=[search_wikipedia,get_current_temperature],
                               verbose=True,
                               memory=memory)
# 调用chain
agent_executor.invoke({"input": "我的名字叫黑大帅"})
agent_executor.invoke({"input": "我是谁"})
